# 示例代码：使用PyTorch加载预训练的ResNet50模型
import io

import torch
from torchvision import models
from torchvision import transforms
from PIL import Image
from scipy.spatial.distance import cosine

model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model.eval()  # 设置为评估模式

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


def compare_features(feat1, feat2):
    return 1 - cosine(torch.flatten(feat1), torch.flatten(feat2))


def extract_img_feature_by_path(img_path):
    img = Image.open(img_path)
    return extract_img_feature(img)


def extract_img_feature(img):
    print(f"Image shape: {img.size}")  # 调试信息
    if img.mode != 'RGB':
        img = img.convert('RGB')  # 确保图像模式为RGB
    img_tensor = preprocess(img).unsqueeze(0)  # 添加批次维度

    with torch.no_grad():  # 禁用梯度计算以节省内存
        features = model(img_tensor)
        print(f"Features shape: {features.shape}")  # 调试信息
        return save_features_to_hex(features)


def save_features_to_hex(features):
    buffer = io.BytesIO()
    torch.save(features, buffer)
    return buffer.getvalue().hex()


def load_features_from_hex(hex_string):
    byte_data = bytes.fromhex(hex_string)
    buffer = io.BytesIO(byte_data)
    return torch.load(buffer)


if __name__ == '__main__':
    # 确保特征向量的维度一致
    feat_a = extract_img_feature_by_path('words/wudo.png')
    feat_b = extract_img_feature_by_path('words/wrong1.png')
    similarity_score = compare_features(load_features_from_hex(feat_a), load_features_from_hex(feat_b))
    print(similarity_score)
